import numpy as np

g = 9.81
R = 10.0
v_sink = 3
valid_window = 20
T = np.array([0, 200, 0])
Of = np.array([0, 0, 0])
M0 = np.array([20000, 0, 2000])
uM = (Of - M0) / np.linalg.norm(Of - M0)
vM = 300 * uM
F0 = np.array([17800, 0, 1800])
v_min = 70
v_max = 140

angles1 = np.deg2rad(np.arange(-180, 180, 22.5))
speeds1 = np.arange(70, 141, 10)
trel1 = np.arange(0, 19, 2)
tdet1 = np.arange(1, 7, 1)
dt1 = 0.01

dt2 = 0.002

best_dur = -np.inf
best_set = []
cnt = 0

for ia, theta in enumerate(angles1):
    ux = np.cos(theta)
    uy = np.sin(theta)

    for iv, vF in enumerate(speeds1):
        for ir, t_rel in enumerate(trel1):
            Dx = F0[0] + vF * ux * t_rel
            Dy = F0[1] + vF * uy * t_rel

            for id, t_det_rel in enumerate(tdet1):
                t_det = t_rel + t_det_rel
                x_det = Dx + vF * ux * t_det_rel
                y_det = Dy + vF * uy * t_det_rel
                z_det = 1800 - 0.5 * g * t_det_rel ** 2

                ts = np.arange(t_det, t_det + valid_window + dt1, dt1)
                if len(ts) == 0:
                    continue

                N = len(ts)
                Mx = M0[0] + vM[0] * ts
                My = M0[1] + vM[1] * ts
                Mz = M0[2] + vM[2] * ts

                Cx = np.full(N, x_det)
                Cy = np.full(N, y_det)
                Cz = z_det - v_sink * (ts - t_det)

                uxL = T[0] - Mx
                uyL = T[1] - My
                uzL = T[2] - Mz
                uu = uxL ** 2 + uyL ** 2 + uzL ** 2

                wx = Cx - Mx
                wy = Cy - My
                wz = Cz - Mz
                wu = wx * uxL + wy * uyL + wz * uzL

                lam = wu / uu
                lam = np.clip(lam, 0, 1)

                Px = Mx + lam * uxL
                Py = My + lam * uyL
                Pz = Mz + lam * uzL

                dx = Cx - Px
                dy = Cy - Py
                dz = Cz - Pz
                d = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2)

                mask = d <= R
                dm = np.diff(mask.astype(int))
                idx_enter = np.where(dm == 1)[0]
                idx_exit = np.where(dm == -1)[0]

                if len(mask) > 0 and mask[0]:
                    idx_enter = np.concatenate(([0], idx_enter))
                if len(mask) > 0 and mask[-1]:
                    idx_exit = np.concatenate((idx_exit, [len(mask) - 1]))

                dur = 0
                m_count = min(len(idx_enter), len(idx_exit))

                for k in range(m_count):
                    i0 = max(0, idx_enter[k])
                    i1 = min(len(ts) - 1, idx_exit[k])

                    t0 = ts[i0]
                    t1 = ts[i0 + 1]
                    g0 = d[i0] - R
                    g1 = d[i0 + 1] - R
                    t_enter_val = t0 - g0 * (t1 - t0) / (g1 - g0) if i0 > 0 else ts[0]

                    t0 = ts[i1]
                    t1 = ts[i1 + 1]
                    g0 = d[i1] - R
                    g1 = d[i1 + 1] - R
                    t_exit_val = t0 - g0 * (t1 - t0) / (g1 - g0) if i1 < len(ts) - 1 else ts[-1]

                    seg = max(0, t_exit_val - t_enter_val)
                    dur += seg

                cnt += 1
                if dur > best_dur:
                    best_dur = dur
                    best_set = [theta, vF, t_rel, t_det_rel]

print(f'阶段1完成：评估了{cnt}个候选。当前最佳遮蔽时长={best_dur:.6f}s')

theta0, v0, t_rel0, t_det_r0 = best_set

angles2 = theta0 + np.deg2rad(np.arange(-10, 11, 2.5))
speeds2 = np.clip(v0 + np.arange(-15, 16, 5), v_min, v_max)
trel2 = np.clip(t_rel0 + np.arange(-0.8, 0.9, 0.2), 0, None)
tdet2 = np.clip(t_det_r0 + np.arange(-0.8, 0.9, 0.2), 0.3, None)

best2_dur = best_dur
best2_set = best_set

for theta in angles2:
    ux = np.cos(theta)
    uy = np.sin(theta)

    for vF in speeds2:
        for t_rel in trel2:
            Dx = F0[0] + vF * ux * t_rel
            Dy = F0[1] + vF * uy * t_rel

            for t_det_rel in tdet2:
                t_det = t_rel + t_det_rel
                x_det = Dx + vF * ux * t_det_rel
                y_det = Dy + vF * uy * t_det_rel
                z_det = 1800 - 0.5 * g * t_det_rel ** 2

                ts = np.arange(t_det, t_det + valid_window + dt2, dt2)
                if len(ts) == 0:
                    continue

                N = len(ts)
                Mx = M0[0] + vM[0] * ts
                My = M0[1] + vM[1] * ts
                Mz = M0[2] + vM[2] * ts

                Cx = np.full(N, x_det)
                Cy = np.full(N, y_det)
                Cz = z_det - v_sink * (ts - t_det)

                uxL = T[0] - Mx
                uyL = T[1] - My
                uzL = T[2] - Mz
                uu = uxL ** 2 + uyL ** 2 + uzL ** 2

                wx = Cx - Mx
                wy = Cy - My
                wz = Cz - Mz
                wu = wx * uxL + wy * uyL + wz * uzL

                lam = wu / uu
                lam = np.clip(lam, 0, 1)

                Px = Mx + lam * uxL
                Py = My + lam * uyL
                Pz = Mz + lam * uzL

                dx = Cx - Px
                dy = Cy - Py
                dz = Cz - Pz
                d = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2)

                mask = d <= R
                dm = np.diff(mask.astype(int))
                idx_enter = np.where(dm == 1)[0]
                idx_exit = np.where(dm == -1)[0]

                if len(mask) > 0 and mask[0]:
                    idx_enter = np.concatenate(([0], idx_enter))
                if len(mask) > 0 and mask[-1]:
                    idx_exit = np.concatenate((idx_exit, [len(mask) - 1]))

                dur = 0
                m_count = min(len(idx_enter), len(idx_exit))

                for k in range(m_count):
                    i0 = max(0, idx_enter[k])
                    i1 = min(len(ts) - 1, idx_exit[k])

                    t0 = ts[i0]
                    t1 = ts[i0 + 1]
                    g0 = d[i0] - R
                    g1 = d[i0 + 1] - R
                    t_enter_val = t0 - g0 * (t1 - t0) / (g1 - g0) if i0 > 0 else ts[0]

                    t0 = ts[i1]
                    t1 = ts[i1 + 1]
                    g0 = d[i1] - R
                    g1 = d[i1 + 1] - R
                    t_exit_val = t0 - g0 * (t1 - t0) / (g1 - g0) if i1 < len(ts) - 1 else ts[-1]

                    seg = max(0, t_exit_val - t_enter_val)
                    dur += seg

                if dur > best2_dur:
                    best2_dur = dur
                    best2_set = [theta, vF, t_rel, t_det_rel]

theta, vF, t_rel, t_det_rel = best2_set
ux = np.cos(theta)
uy = np.sin(theta)
Dx = F0[0] + vF * ux * t_rel
Dy = F0[1] + vF * uy * t_rel
Dz = 1800
t_det = t_rel + t_det_rel
x_det = Dx + vF * ux * t_det_rel
y_det = Dy + vF * uy * t_det_rel
z_det = 1800 - 0.5 * g * t_det_rel ** 2

print('\n========= 问题2最优单弹遮蔽策略（FY1→M1） =========')
print(f'航向角θ = {np.rad2deg(theta):.2f}°（相对+x逆时针）')
print(f'飞行速度 v = {vF:.2f} m/s')
print(f'投放时刻 t_rel = {t_rel:.3f} s')
print(f'引信延时 t_det_rel = {t_det_rel:.3f} s')
print(f'投放点 D = ({Dx:.3f}, {Dy:.3f}, {Dz:.3f}) m')
print(f'起爆点 C0 = ({x_det:.3f}, {y_det:.3f}, {z_det:.3f}) m')
print(f'起爆全局时刻 t_det = {t_det:.3f} s')
print(f'最大遮蔽总时长 = {best2_dur:.6f} s')